package com.webbdong.aop.agent.test.bytebuddy;

import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.modifier.Ownership;
import net.bytebuddy.description.modifier.Visibility;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.dynamic.loading.ClassInjector;
import net.bytebuddy.implementation.LoadedTypeInitializer;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.implementation.bind.annotation.AllArguments;
import net.bytebuddy.implementation.bind.annotation.Origin;
import net.bytebuddy.implementation.bind.annotation.RuntimeType;
import net.bytebuddy.implementation.bind.annotation.SuperCall;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.pool.TypePool;
import net.bytebuddy.utility.JavaModule;

import java.io.File;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.nio.file.Files;
import java.util.Collections;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @program: spring-agent
 * @description:
 * @author: zt22328 赵儒定
 * @create: 2023-04-04 14:56
 **/
public class BootstrapAgent {

    public static void main(String[] args) throws Exception {
        premain(null, ByteBuddyAgent.install());
        final ExecutorService executorService = Executors.newCachedThreadPool();
    }

    public static void premain(String arg, Instrumentation instrumentation) throws Exception {
        File temp = Files.createTempDirectory("tmp").toFile();
        ClassInjector.UsingInstrumentation.of(temp, ClassInjector.UsingInstrumentation.Target.BOOTSTRAP, instrumentation)
                .inject(Collections.singletonMap(
                        new TypeDescription.ForLoadedType(GeneralInterceptor.class),
                        ClassFileLocator.ForClassLoader.read(GeneralInterceptor.class)));

        String interceptor = "GeneralInterceptor";
        ClassFileLocator locator = ClassFileLocator.ForClassLoader.ofSystemLoader();
        TypePool pool = TypePool.Default.of(locator);
        final TypeDescription target = pool.describe(interceptor).resolve();

//        File temp = Files.createTempDirectory("tmp").toFile();
        ClassInjector.UsingInstrumentation
                .of(temp, ClassInjector.UsingInstrumentation.Target.BOOTSTRAP, instrumentation)
                .injectRaw(Collections.singletonMap(interceptor, locator.locate(interceptor).resolve()));
        new AgentBuilder.Default()
                .ignore(ElementMatchers.nameStartsWith("net.bytebuddy."))
                .with(new AgentBuilder.InjectionStrategy.UsingInstrumentation(instrumentation, temp))
                .type(ElementMatchers.nameEndsWith("java.util.concurrent.ThreadPoolExecutor"))
                .transform((DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule javaModule) -> {
                        return builder.defineField("delegator", target, Ownership.STATIC, Visibility.PRIVATE)
                                .initializer(new LoadedTypeInitializer() {
                                    @Override
                                    public void onLoad(Class<?> type) {
                                        try {
                                            Field field = type.getDeclaredField("delegator");
                                            field.setAccessible(true);
                                            field.set(null, Class.forName(interceptor, false, null).getConstructor().newInstance());
                                        } catch (Exception e) {
                                            throw new RuntimeException(e);
                                        }
                                    }

                                    @Override
                                    public boolean isAlive() {
                                        return true;
                                    }
                                })
                                // Interception constructor not work
                                // .method(ElementMatchers.named("runWorker")) is work
                                .method(ElementMatchers.isConstructor())
                                .intercept(MethodDelegation.toField("delegator"));
                    })
                .installOn(instrumentation);
    }

}

